library(dplyr)
library(stringr)
library(ggplot2)
library(scales)
library(ggrepel)
library(ggpubr)

if (!file.exists("output")) {
  dir.create("output")
}

# NOTE: Output obtained from risk_factor_regressions.R, also shown in manuscript Supplementary Table 13
LR <- read.csv("ExtendedDataFigure_5ab_data.csv") %>%
  mutate(
    p_adj_log = -log10(p_adj),
    sig_type = case_when(
      grepl("SBS", signature) ~ "SBS",
      grepl("DBS", signature) ~ "DBS",
      grepl("ID", signature) ~ "ID"
    )
  )

# variables to plot
variables <- c("alcohol yes", "tobacco yes")

palette <- c("#1b9e77", "#d95f02", "#7570b3")
names(palette) <- unique(LR$sig_type)
threshold <- 2

for (var in variables) {
  LR.var <- LR[LR$Independent_variable == var, ]

  var_simple <- str_split_i(var, " ", i = 1)
  correct <- c("sex", "age of diagnosis", "subsite", "region", "tobacco", "alcohol")
  correct.minus <- correct[!grepl(var_simple, correct)]
  subtitle <- paste0(
    "Corrected for ",
    paste(correct.minus[-length(correct.minus)], collapse = ", "),
    ", and ", correct.minus[length(correct.minus)]
  )

  p <- ggplot(LR.var, aes(
    x = OR, y = p_adj_log,
    color = ifelse(
      (OR > threshold | OR < 1 / threshold) & p_adj < 0.05, sig_type, NA
    )
  )) +
    geom_point(
      stroke = 0, size = 2.4,
      alpha = ifelse((LR.var$OR > threshold | LR.var$OR < 1 / threshold) & LR.var$p_adj < 0.05, 1, 0.6)
    ) +
    geom_text_repel(
      label = ifelse((LR.var$OR > threshold | LR.var$OR < 1 / threshold) & LR.var$p_adj < 0.05, LR.var$signature, NA),
      show.legend = FALSE
    ) +
    geom_hline(yintercept = -log10(0.05), linetype = "dashed", color = "gray") +
    geom_vline(xintercept = threshold, linetype = "dashed", color = "gray") +
    geom_vline(xintercept = 1 / threshold, linetype = "dashed", color = "gray") +
    scale_x_continuous(trans = log2_trans()) +
    scale_y_continuous(limits = c(0, 8)) +
    scale_color_manual(values = palette, breaks = unique(LR.var$sig_type)) +
    labs(
      x = "OR", y = "Adjusted p (-log 10)", color = "",
      title = paste("Associations with", var_simple),
      subtitle = subtitle
    ) +
    theme_minimal() +
    theme(
      legend.position = "right",
      legend.text = element_text(size = 12),
      plot.title = element_text(face = "bold", size = 12),
      plot.subtitle = element_text(size = 9),
      axis.text = element_text(size = 12),
      axis.line = element_line(size = .3,color = "#404040"),
      axis.ticks = element_line(size = .3,color = "#404040"),
      panel.grid = element_blank()
    )
  print(p)
  ggsave(
    filename = paste0("output/ExtendedDataFigure_5ab_", str_replace_all(var, " ", "_"), "_", Sys.Date(), ".pdf"),
    device = "pdf",
    width = 5, height = 4, dpi = 700
  )
}
